{
    "cells": [
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "# Solve Burgers' equation based on Fourier Neural Operator\n",
                "\n",
                "[![DownloadNotebook](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_notebook_en.png)](https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/master/mindflow/en/data_driven/mindspore_burgers_FNO1D.ipynb)&emsp;[![DownloadCode](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_download_code_en.png)](https://obs.dualstack.cn-north-4.myhuaweicloud.com/mindspore-website/notebook/master/mindflow/en/data_driven/mindspore_burgers_FNO1D.py)&emsp;[![ViewSource](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/master/resource/_static/logo_source_en.png)](https://gitee.com/mindspore/docs/blob/master/docs/mindflow/docs/source_en/data_driven/burgers_FNO1D.ipynb)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "## Overview\n",
                "\n",
                "Computational fluid dynamics is one of the most important techniques in the field of fluid mechanics in the 21st century. The flow analysis, prediction and control can be realized by solving the governing equations of fluid mechanics by numerical method. Traditional finite element method (FEM) and finite difference method (FDM) are inefficient because of the complex simulation process (physical modeling, meshing, numerical discretization, iterative solution, etc.) and high computing costs. Therefore, it is necessary to improve the efficiency of fluid simulation with AI.\n",
                "\n",
                "Machine learning methods provide a new paradigm for scientific computing by providing a fast solver similar to traditional methods. Classical neural networks learn mappings between finite dimensional spaces and can only learn solutions related to a specific discretization. Different from traditional neural networks, Fourier Neural Operator (FNO) is a new deep learning architecture that can learn mappings between infinite-dimensional function spaces. It directly learns mappings from arbitrary function parameters to solutions to solve a class of partial differential equations.  Therefore, it has a stronger generalization capability. More information can be found in the paper, [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/abs/2010.08895).\n",
                "\n",
                "This tutorial describes how to solve the 1-d Burgers' equation using Fourier neural operator."
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "## Burgers' equation\n",
                "\n",
                "The 1-d Burgers’ equation is a non-linear PDE with various applications including modeling the one\n",
                "dimensional flow of a viscous fluid. It takes the form\n",
                "\n",
                "$$\n",
                "\\partial_t u(x, t)+\\partial_x (u^2(x, t)/2)=\\nu \\partial_{xx} u(x, t), \\quad x \\in(0,1), t \\in(0, 1]\n",
                "$$\n",
                "\n",
                "$$\n",
                "u(x, 0)=u_0(x), \\quad x \\in(0,1)\n",
                "$$\n",
                "\n",
                "where $u$ is the velocity field, $u_0$ is the initial condition and $\\nu$ is the viscosity coefficient.\n"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Problem Description\n",
                "\n",
                "We aim to learn the operator mapping the initial condition to the solution at time one:\n",
                "\n",
                "$$\n",
                "u_0 \\mapsto u(\\cdot, 1)\n",
                "$$"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Technology Path\n",
                "\n",
                "MindFlow solves the problem as follows:\n",
                "\n",
                "1. Training Dataset Construction.\n",
                "2. Model Construction.\n",
                "3. Optimizer and Loss Function.\n",
                "4. Model Training."
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "## Fourier Neural Operator\n",
                "\n",
                "The Fourier Neural Operator consists of the Lifting Layer, Fourier Layers, and the Decoding Layer.\n",
                "\n",
                "![Fourier Neural Operator model structure](images/FNO.png)\n",
                "\n",
                "Fourier layers: Start from input V. On top: apply the Fourier transform $\\mathcal{F}$; a linear transform R on the lower Fourier modes and filters out the higher modes; then apply the inverse Fourier transform $\\mathcal{F}^{-1}$. On the bottom: apply a local linear transform W.  Finally, the Fourier Layer output vector is obtained through the activation function.\n",
                "\n",
                "![Fourier Layer structure](images/FNO-2.png)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 1,
            "metadata": {
                "pycharm": {
                    "name": "#%%\n"
                }
            },
            "outputs": [],
            "source": [
                "import os\n",
                "import time\n",
                "import numpy as np\n",
                "\n",
                "from mindspore.amp import DynamicLossScaler, auto_mixed_precision, all_finite\n",
                "from mindspore import context, nn, Tensor, set_seed, ops, data_sink, jit, save_checkpoint\n",
                "from mindspore import dtype as mstype\n",
                "from mindflow import FNO1D, RelativeRMSELoss, load_yaml_config, get_warmup_cosine_annealing_lr\n",
                "from mindflow.pde import UnsteadyFlowWithLoss"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "The following `src` pacakage can be downloaded in [applications/data_driven/burgers/fno1d/src](https://gitee.com/mindspore/mindscience/tree/master/MindFlow/applications/data_driven/burgers/fno1d/src)."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 2,
            "metadata": {
                "pycharm": {
                    "name": "#%%\n"
                }
            },
            "outputs": [],
            "source": [
                "from src import create_training_dataset\n",
                "\n",
                "set_seed(0)\n",
                "np.random.seed(0)\n",
                "\n",
                "context.set_context(mode=context.GRAPH_MODE, device_target=\"GPU\", device_id=0)\n",
                "use_ascend = context.get_context(attr_key='device_target') == \"Ascend\""
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "You can get parameters of model, data and optimizer from `config`."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 3,
            "metadata": {
                "pycharm": {
                    "name": "#%%\n"
                }
            },
            "outputs": [],
            "source": [
                "config = load_yaml_config('./configs/fno1d.yaml')\n",
                "data_params = config[\"data\"]\n",
                "model_params = config[\"model\"]\n",
                "optimizer_params = config[\"optimizer\"]"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "## Training Dataset Construction\n",
                "\n",
                "Download the training and test dataset: [data_driven/burgers/dataset](https://download.mindspore.cn/mindscience/mindflow/dataset/applications/data_driven/burgers/dataset/) .\n",
                "\n",
                "In this case, training datasets and test datasets are generated according to Zongyi Li's dataset in [Fourier Neural Operator for Parametric Partial Differential Equations](https://arxiv.org/pdf/2010.08895.pdf) . The settings are as follows:\n",
                "\n",
                "the initial condition $u_0(x)$ is generated according to periodic boundary conditions:\n",
                "\n",
                "$$\n",
                "u_0 \\sim \\mu, \\mu=\\mathcal{N}\\left(0,625(-\\Delta+25 I)^{-2}\\right)\n",
                "$$\n",
                "\n",
                "We set the viscosity to $\\nu=0.1$ and solve the equation using a split step method where the heat equation part is solved exactly in Fourier space then the non-linear part is advanced, again in Fourier space, using a very fine forward Euler method. The number of samples in the training set is 1000, and the number of samples in the test set is 200.\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 4,
            "metadata": {
                "pycharm": {
                    "name": "#%%\n"
                }
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "Data preparation finished\n",
                        "input_path:  (1000, 1024, 1)\n",
                        "label_path:  (1000, 1024)\n"
                    ]
                }
            ],
            "source": [
                "# create training dataset\n",
                "train_dataset = create_training_dataset(data_params, model_params, shuffle=True)\n",
                "\n",
                "# create test dataset\n",
                "test_input, test_label = np.load(os.path.join(data_params[\"root_dir\"], \"test/inputs.npy\")), \\\n",
                "                         np.load(os.path.join(data_params[\"root_dir\"], \"test/label.npy\"))\n",
                "test_input = Tensor(np.expand_dims(test_input, -2), mstype.float32)\n",
                "test_label = Tensor(np.expand_dims(test_label, -2), mstype.float32)"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "## Model Construction\n",
                "\n",
                "The network is composed of 1 lifting layer, multiple Fourier layers and 1 decoding layer:\n",
                "\n",
                "- The Lifting layer corresponds to the `FNO1D.fc0` in the case, and maps the output data $x$ to the high dimension;\n",
                "\n",
                "- Multi-layer Fourier Layer corresponds to the `FNO1D.fno_seq` in the case. Discrete Fourier transform is used to realize the conversion between time domain and frequency domain;\n",
                "\n",
                "- The Decoding layer corresponds to `FNO1D.fc1` and `FNO1D.fc2` in the case to obtain the final predictive value.\n",
                "\n",
                "The initialization of the model based on the network above, parameters can be modified in [configuration file](https://gitee.com/mindspore/mindscience/blob/master/MindFlow/applications/data_driven/burgers/fno1d/configs/fno1d.yaml)."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 5,
            "metadata": {
                "pycharm": {
                    "name": "#%%\n"
                }
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "name:FNO1D_in_channels:1_out_channels:1_resolution:1024_modes:16_width:64_depth:4\n"
                    ]
                }
            ],
            "source": [
                "model = FNO1D(in_channels=model_params[\"in_channels\"],\n",
                "              out_channels=model_params[\"out_channels\"],\n",
                "              n_modes=model_params[\"modes\"],\n",
                "              resolutions=model_params[\"resolutions\"],\n",
                "              hidden_channels=model_params[\"hidden_channels\"],\n",
                "              n_layers=model_params[\"depths\"],\n",
                "              projection_channels=4*model_params[\"hidden_channels\"],\n",
                "              )\n",
                "\n",
                "model_params_list = []\n",
                "for k, v in model_params.items():\n",
                "    model_params_list.append(f\"{k}:{v}\")\n",
                "model_name = \"_\".join(model_params_list)\n",
                "print(model_name)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "## Optimizer and Loss Function"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 6,
            "metadata": {
                "pycharm": {
                    "name": "#%%\n"
                }
            },
            "outputs": [],
            "source": [
                "steps_per_epoch = train_dataset.get_dataset_size()\n",
                "lr = get_warmup_cosine_annealing_lr(lr_init=optimizer_params[\"learning_rate\"],\n",
                "                                    last_epoch=optimizer_params[\"epochs\"],\n",
                "                                    steps_per_epoch=steps_per_epoch,\n",
                "                                    warmup_epochs=1)\n",
                "optimizer = nn.Adam(model.trainable_params(), learning_rate=Tensor(lr))\n",
                "\n",
                "if use_ascend:\n",
                "    from mindspore.amp import DynamicLossScaler, auto_mixed_precision, all_finite\n",
                "    loss_scaler = DynamicLossScaler(1024, 2, 100)\n",
                "    auto_mixed_precision(model, 'O1')\n",
                "else:\n",
                "    loss_scaler = None"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {
                "pycharm": {
                    "name": "#%% md\n"
                }
            },
            "source": [
                "## Model Training\n",
                "\n",
                "With **MindSpore version >= 2.0.0**, we can use the functional programming for training neural networks. `MindFlow` provide a training interface for unsteady problems `UnsteadyFlowWithLoss` for model training and evaluation."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": 7,
            "metadata": {
                "pycharm": {
                    "name": "#%%\n"
                }
            },
            "outputs": [
                {
                    "name": "stdout",
                    "output_type": "stream",
                    "text": [
                        "pid: 1993\n",
                        "2023-02-01 12:14:12.2323\n",
                        "use_ascend: False\n",
                        "device_id: 4\n",
                        "Data preparation finished\n",
                        "input_path:  (1000, 1024, 1)\n",
                        "label_path:  (1000, 1024)\n",
                        "name:FNO1D_in_channels:1_out_channels:1_resolution:1024_modes:16_width:64_depth:4\n",
                        "./summary/name:FNO1D_in_channels:1_out_channels:1_resolution:1024_modes:16_width:64_depth:4\n",
                        "epoch: 1 train loss: 2.167046070098877 epoch time: 21.75s\n",
                        "epoch: 2 train loss: 0.5935954451560974 epoch time: 5.53s\n",
                        "epoch: 3 train loss: 0.7349425554275513 epoch time: 5.46s\n",
                        "epoch: 4 train loss: 0.6338694095611572 epoch time: 4.95s\n",
                        "epoch: 5 train loss: 0.3174982964992523 epoch time: 5.57s\n",
                        "epoch: 6 train loss: 0.3099440038204193 epoch time: 5.71s\n",
                        "epoch: 7 train loss: 0.3117891848087311 epoch time: 5.22s\n",
                        "epoch: 8 train loss: 0.1810857653617859 epoch time: 4.82s\n",
                        "epoch: 9 train loss: 0.1386510729789734 epoch time: 4.97s\n",
                        "epoch: 10 train loss: 0.2102256715297699 epoch time: 4.85s\n",
                        "mean rms_error: 0.027940063\n",
                        "=================================End Evaluation=================================\n",
                        "...\n",
                        "epoch: 91 train loss: 0.019643772393465042 epoch time: 4.40s\n",
                        "epoch: 92 train loss: 0.0641067773103714 epoch time: 5.48s\n",
                        "epoch: 93 train loss: 0.02199840545654297 epoch time: 5.55s\n",
                        "epoch: 94 train loss: 0.024467874318361282 epoch time: 6.24s\n",
                        "epoch: 95 train loss: 0.025712188333272934 epoch time: 5.43s\n",
                        "epoch: 96 train loss: 0.02247200347483158 epoch time: 6.48s\n",
                        "epoch: 97 train loss: 0.026637140661478043 epoch time: 6.30s\n",
                        "epoch: 98 train loss: 0.030040305107831955 epoch time: 5.16s\n",
                        "epoch: 99 train loss: 0.02589748054742813 epoch time: 5.36s\n",
                        "epoch: 100 train loss: 0.028599221259355545 epoch time: 5.90s\n",
                        "================================Start Evaluation================================\n",
                        "mean rms_error: 0.0037017763\n",
                        "=================================End Evaluation=================================\n"
                    ]
                }
            ],
            "source": [
                "problem = UnsteadyFlowWithLoss(model, loss_fn=RelativeRMSELoss(), data_format=\"NHWTC\")\n",
                "\n",
                "summary_dir = os.path.join(config[\"summary\"][\"summary_dir\"], model_name)\n",
                "print(summary_dir)\n",
                "\n",
                "def forward_fn(data, label):\n",
                "    loss = problem.get_loss(data, label)\n",
                "    if use_ascend:\n",
                "        loss = loss_scaler.scale(loss)\n",
                "    return loss\n",
                "\n",
                "grad_fn = ops.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=False)\n",
                "\n",
                "@jit\n",
                "def train_step(data, label):\n",
                "    loss, grads = grad_fn(data, label)\n",
                "    if use_ascend:\n",
                "        loss = loss_scaler.unscale(loss)\n",
                "        if all_finite(grads):\n",
                "            grads = loss_scaler.unscale(grads)\n",
                "    loss = ops.depend(loss, optimizer(grads))\n",
                "    return loss\n",
                "\n",
                "sink_process = data_sink(train_step, train_dataset, 1)\n",
                "summary_dir = os.path.join(config[\"summary_dir\"], model_name)\n",
                "ckpt_dir = os.path.join(summary_dir, \"ckpt\")\n",
                "if not os.path.exists(ckpt_dir):\n",
                "    os.makedirs(ckpt_dir)\n",
                "\n",
                "for epoch in range(1, optimizer_params[\"epochs\"] + 1):\n",
                "    model.set_train()\n",
                "    local_time_beg = time.time()\n",
                "    for _ in range(steps_per_epoch):\n",
                "        cur_loss = sink_process()\n",
                "    print(\n",
                "        f\"epoch: {epoch} train loss: {cur_loss.asnumpy():.8f}\"\\\n",
                "        f\" epoch time: {time.time() - local_time_beg:.2f}s\"\\\n",
                "        f\" step time: {(time.time() - local_time_beg)/steps_per_epoch:.4f}s\")\n",
                "\n",
                "    if epoch % config['summary']['test_interval'] == 0:\n",
                "        model.set_train(False)\n",
                "        print(\"================================Start Evaluation================================\")\n",
                "        rms_error = problem.get_loss(test_input, test_label)/test_input.shape[0]\n",
                "        print(f\"mean rms_error: {rms_error}\")\n",
                "        print(\"=================================End Evaluation=================================\")\n",
                "        save_checkpoint(model, os.path.join(ckpt_dir, model_params[\"name\"] + '_epoch' + str(epoch)))"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.10.6"
        },
        "vscode": {
            "interpreter": {
                "hash": "916dbcbb3f70747c44a77c7bcd40155683ae19c65e1c03b4aa3499c5328201f1"
            }
        }
    },
    "nbformat": 4,
    "nbformat_minor": 1
}
